/*
 * Unidata Platform Community Edition
 * Copyright (c) 2013-2020, UNIDATA LLC, All rights reserved.
 * This file is part of the Unidata Platform Community Edition software.
 *
 * Unidata Platform Community Edition is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Unidata Platform Community Edition is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <https://www.gnu.org/licenses/>.
 */
package org.unidata.mdm.dq.core.service.impl;

import java.time.OffsetDateTime;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.dao.DuplicateKeyException;
import org.springframework.stereotype.Component;
import org.unidata.mdm.core.context.ModelChangeContext;
import org.unidata.mdm.core.context.ModelChangeContext.ModelChangeType;
import org.unidata.mdm.core.context.ModelGetContext;
import org.unidata.mdm.core.context.ModelRefreshContext;
import org.unidata.mdm.core.context.ModelRemoveContext;
import org.unidata.mdm.core.context.ModelSourceContext;
import org.unidata.mdm.core.dto.ModelGetResult;
import org.unidata.mdm.core.service.MetaModelService;
import org.unidata.mdm.core.service.UPathService;
import org.unidata.mdm.core.service.impl.AbstractModelComponent;
import org.unidata.mdm.core.type.model.ModelDescriptor;
import org.unidata.mdm.core.type.model.ModelImplementation;
import org.unidata.mdm.core.type.model.ModelSource;
import org.unidata.mdm.core.util.SecurityUtils;
import org.unidata.mdm.dq.core.configuration.DataQualityDescriptors;
import org.unidata.mdm.dq.core.configuration.DataQualityModelIds;
import org.unidata.mdm.dq.core.context.GetQualityModelContext;
import org.unidata.mdm.dq.core.context.SourceQualityModelContext;
import org.unidata.mdm.dq.core.context.UpsertQualityModelContext;
import org.unidata.mdm.dq.core.dao.DataQualityDAO;
import org.unidata.mdm.dq.core.dto.GetAssignmentsModelResult;
import org.unidata.mdm.dq.core.dto.GetFunctionsModelResult;
import org.unidata.mdm.dq.core.dto.GetGroupsModelResult;
import org.unidata.mdm.dq.core.dto.GetQualityModelResult;
import org.unidata.mdm.dq.core.dto.GetRulesModelResult;
import org.unidata.mdm.dq.core.dto.GetSetsModelResult;
import org.unidata.mdm.dq.core.exception.DataQualityExceptionIds;
import org.unidata.mdm.dq.core.exception.DataQualityRuntimeException;
import org.unidata.mdm.dq.core.exception.DataQualityValidationException;
import org.unidata.mdm.dq.core.po.DataQualityPO;
import org.unidata.mdm.dq.core.serialization.DataQualitySerializer;
import org.unidata.mdm.dq.core.service.impl.instance.AbstractCleanseFunctionImpl;
import org.unidata.mdm.dq.core.service.impl.instance.DataQualityInstanceImpl;
import org.unidata.mdm.dq.core.service.impl.instance.FunctionGroupImpl;
import org.unidata.mdm.dq.core.service.impl.instance.MappingSetImpl;
import org.unidata.mdm.dq.core.service.impl.instance.NamespaceAssignmentImpl;
import org.unidata.mdm.dq.core.service.impl.instance.QualityRuleImpl;
import org.unidata.mdm.dq.core.type.draft.DataQualityDraftConstants;
import org.unidata.mdm.dq.core.type.model.instance.DataQualityInstance;
import org.unidata.mdm.dq.core.type.model.source.AbstractCleanseFunctionSource;
import org.unidata.mdm.dq.core.type.model.source.CleanseFunctionGroup;
import org.unidata.mdm.dq.core.type.model.source.CompositeCleanseFunctionSource;
import org.unidata.mdm.dq.core.type.model.source.DataQualityModel;
import org.unidata.mdm.dq.core.type.model.source.GroovyCleanseFunctionSource;
import org.unidata.mdm.dq.core.type.model.source.JavaCleanseFunctionSource;
import org.unidata.mdm.dq.core.type.model.source.PythonCleanseFunctionSource;
import org.unidata.mdm.dq.core.type.model.source.assignment.NameSpaceAssignmentSource;
import org.unidata.mdm.dq.core.type.model.source.rule.MappingSetSource;
import org.unidata.mdm.dq.core.type.model.source.rule.QualityRuleSource;
import org.unidata.mdm.dq.core.util.DQUtils;
import org.unidata.mdm.draft.context.DraftGetContext;
import org.unidata.mdm.draft.context.DraftUpsertContext;
import org.unidata.mdm.draft.dto.DraftGetResult;
import org.unidata.mdm.draft.service.DraftService;
import org.unidata.mdm.draft.type.Draft;
import org.unidata.mdm.draft.type.Edition;
import org.unidata.mdm.system.exception.PlatformValidationException;
import org.unidata.mdm.system.exception.ValidationResult;
import org.unidata.mdm.system.service.AfterModuleStartup;

/**
 * @author Mikhail Mikhailov on Jan 21, 2021
 * <br>
 * Data Quality model implementation.
 */
@Component(DataQualityModelIds.DATA_QUALITY)
public class DataQualityModelComponent extends AbstractModelComponent implements ModelImplementation<DataQualityInstance>, AfterModuleStartup {
    /**
     * The MMS.
     */
    private MetaModelService metaModelService;
    /**
     * The DS.
     */
    @Autowired
    private DraftService draftService;
    /**
     * UPS.
     */
    @Autowired
    private UPathService upathService;
    /**
     * The CFCC.
     */
    @Autowired
    private CleanseFunctionCacheComponent cleanseFunctionCacheComponent;
    /**
     * DQ VC.
     */
    @Autowired
    private DataQualityValidationComponent dataQualityValidationComponent;
    /**
     * DQ DAO.
     */
    @Autowired
    private DataQualityDAO dataQualityDAO;
    /**
     * Model instances, keyed by storage ID.
     */
    private final Map<String, DataQualityInstance> instances = new ConcurrentHashMap<>(4);
    /**
     * Constructor.
     */
    public DataQualityModelComponent(MetaModelService metaModelService) {
        super();
        this.metaModelService = metaModelService;
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public void afterModuleStartup() {
        metaModelService.register(this);
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public ModelDescriptor<DataQualityInstance> descriptor() {
        return DataQualityDescriptors.DQ;
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public DataQualityInstance instance(String storageId, String id) {
        return instances.computeIfAbsent(storageId, this::load);
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public DataQualityInstance instance(Long draftId, String storageId, String id) {

        if (Objects.nonNull(draftId) && draftId.longValue() > 0) {

            Edition current = draftService.current(draftId, true);
            DataQualityModel m =  Objects.isNull(current)
                    ? new DataQualityModel().withVersion(0)
                    : current.getContent();

            return new DataQualityInstanceImpl(m, cleanseFunctionCacheComponent, upathService);
        }

        return instance(storageId, id);
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public ModelGetResult get(ModelGetContext ctx) {

        GetQualityModelContext get = (GetQualityModelContext) ctx;

        if (get.isDraftOperation()) {

            Long draftId = get.getDraftId();
            Long parentDraftId = get.getParentDraftId();

            DraftGetResult result = draftService.get(DraftGetContext.builder()
                    .draftId(draftId)
                    .parentDraftId(parentDraftId)
                    .payload(get)
                    .build());

            return result.hasPayload() ? result.narrow() : null;
        }

        return disassemble(instance(SecurityUtils.getStorageId(get), null), get);
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public void upsert(ModelChangeContext ctx) {

        UpsertQualityModelContext change = (UpsertQualityModelContext) ctx;

        // 1. Process draft and return
        if (change.isDraftOperation()) {

            Long draftId = change.getDraftId();
            Long parentDraftId = change.getParentDraftId();

            draftService.upsert(DraftUpsertContext.builder()
                    .draftId(draftId)
                    .parentDraftId(parentDraftId)
                    .payload(change)
                    .build());

            return;
        }

        // 2. Run direct upsert
        DataQualityModel target = assemble(change);

        List<ValidationResult> validations = new ArrayList<>();
        validations.addAll(validate(target));
        validations.addAll(metaModelService.allow(SourceQualityModelContext.builder()
                    .source(target)
                    .build()));

        if (CollectionUtils.isNotEmpty(validations)) {
            throw new DataQualityValidationException("Cannot upsert data quality model. Validation errors exist.",
                    DataQualityExceptionIds.EX_DQ_UPSERT_DATA_MODEL_VALIDATION, validations);
        }

        put(SecurityUtils.getStorageId(change), target, change.force());
    }
    /**
     * Puts a model to storage.
     * @param storageId the target storage id
     * @param src the model source
     * @param force force to latest version upon failure
     */
    public void put(String storageId, ModelSource src, boolean force) {

        Objects.requireNonNull(storageId, "Storage id must not be null.");
        DataQualityModel source = (DataQualityModel) src;

        DataQualityPO po = new DataQualityPO();
        po.setCreatedBy(SecurityUtils.getCurrentUserName());
        po.setStorageId(storageId);
        po.setRevision(source.getVersion());
        po.setContent(DataQualitySerializer.modelToCompressedXml(source));

        boolean running = true;
        while (running) {

            try {
                dataQualityDAO.save(po);
            } catch (DuplicateKeyException dke) {

                if (force) {

                    int latest = dataQualityDAO.latest(storageId);
                    int next = latest + 1;

                    source.withVersion(next);

                    po.setRevision(next);
                    po.setContent(DataQualitySerializer.modelToCompressedXml(source));

                    continue;
                } else {
                    throw new DataQualityRuntimeException(
                            "Cannot save data model. Revisions conflict [expected next {}].",
                            dke,
                            DataQualityExceptionIds.EX_DQ_UPSERT_DATA_MODEL_REVISION_EXISTS,
                            po.getRevision());
                }
            }

            running = false;
        }
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public DataQualityModel assemble(ModelChangeContext ctx) {

        UpsertQualityModelContext change = (UpsertQualityModelContext) ctx;
        String storageId = SecurityUtils.getStorageId(change);

        DataQualityInstance current = instance(storageId, null);

        DataQualityModel source;
        DataQualityModel target;
        Integer next;

        // 1. Processing draft
        if (change.isDraftOperation()) {

            // Must be set!
            Draft draft = change.currentDraft();

            // 1.2. Load last edition
            Edition latest = null;
            if (draft.isExisting()) {
                latest = draftService.current(draft.getDraftId(), true);
            }

            // 1.3. Take data snapshot for new drafts, if no editions exist
            if (Objects.isNull(latest)) {
                source = current.toSource();
                next = current.getVersion() + 1;
            // 1.4. Use previous revision data otherwise
            } else {
                source = latest.getContent();
                next = draft.getVariables().<Integer>valueGet(DataQualityDraftConstants.DRAFT_START_VERSION) + 1;
            }

        // 2. Direct upsert
        } else {
            source = change.getUpsertType() == ModelChangeType.FULL ? null : current.toSource();
            next = Objects.nonNull(change.getVersion()) ? change.getVersion() : current.getVersion() + 1;
        }

        target = new DataQualityModel()
                .withVersion(next)
                .withStorageId(storageId)
                .withCreateDate(OffsetDateTime.now())
                .withCreatedBy(SecurityUtils.getCurrentUserName());

        // 3. Process
        if (change.getUpsertType() == ModelChangeType.FULL) {
            target
              .withCompositeFunctions(change.getCompositeFunctionsUpdate())
              .withGroovyFunctions(change.getGroovyFunctionsUpdate())
              .withJavaFunctions(change.getJavaFunctionsUpdate())
              .withPythonFunctions(change.getPythonFunctionsUpdate())
              .withFunctionGroup(change.getGroupsUpdate())
              .withRules(change.getRulesUpdate())
              .withSets(change.getSetsUpdate())
              .withAssignments(change.getAssignmentsUpdate());
        } else {
            merge(target, source, change);
        }

        // 4. Info fields
        processInfoFields(target, change);

        // 5. Pre-check before validation to detect serious violations, rendering model unusable.
        Collection<ValidationResult> validations = dataQualityValidationComponent.precheck(target);
        if (CollectionUtils.isNotEmpty(validations)) {
            throw new PlatformValidationException("Cannot save data quality model. Pre-check validation errors exist.",
                    DataQualityExceptionIds.EX_DQ_UPSERT_DATA_MODEL_CONSISTENCY, validations);
        }

        return target;
    }

    public GetQualityModelResult disassemble(DataQualityInstance i, GetQualityModelContext ctx) {

        GetQualityModelResult result = new GetQualityModelResult();

        processFunctions(i, ctx, result);
        processFunctionGroups(i, ctx, result);
        processRules(i, ctx, result);
        processSets(i, ctx, result);
        processAssignments(i, ctx, result);
        processInfoFields(i, ctx, result);

        return result;
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public Collection<ValidationResult> validate(ModelSource source) {
        return dataQualityValidationComponent.validate((DataQualityModel) source);
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public Collection<ValidationResult> allow(ModelSourceContext<?> source) {
        return dataQualityValidationComponent.allow(source);
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public void refresh(ModelRefreshContext refresh) {
        instances.computeIfPresent(SecurityUtils.getStorageId(refresh), (k, v) -> load(k));
    }
    /**
     * {@inheritDoc}
     */
    @Override
    public void remove(ModelRemoveContext remove) {
        // Not supported
    }

    private void merge(DataQualityModel target, DataQualityModel source, UpsertQualityModelContext change) {

        // 1. Source
        Map<String, JavaCleanseFunctionSource> javaFunctionsMerge = source.getJavaFunctions().stream()
                .collect(Collectors.toMap(AbstractCleanseFunctionSource::getName, Function.identity()));

        Map<String, CompositeCleanseFunctionSource> compositeFunctionsMerge = source.getCompositeFunctions().stream()
                .collect(Collectors.toMap(AbstractCleanseFunctionSource::getName, Function.identity()));

        Map<String, PythonCleanseFunctionSource> pythonFunctionsMerge = source.getPythonFunctions().stream()
                .collect(Collectors.toMap(AbstractCleanseFunctionSource::getName, Function.identity()));

        Map<String, GroovyCleanseFunctionSource> groovyFunctionsMerge = source.getGroovyFunctions().stream()
                .collect(Collectors.toMap(AbstractCleanseFunctionSource::getName, Function.identity()));

        Map<String, QualityRuleSource> rulesMerge = source.getRules().stream()
                .collect(Collectors.toMap(QualityRuleSource::getName, Function.identity()));

        Map<String, MappingSetSource> setsMerge = source.getSets().stream()
                .collect(Collectors.toMap(MappingSetSource::getName, Function.identity()));

        Map<String, NameSpaceAssignmentSource> assignmentsMerge = source.getAssignments().stream()
                .collect(Collectors.toMap(NameSpaceAssignmentSource::getNameSpace, Function.identity()));

        // 2. Updates
        if (change.hasJavaFunctionsUpdate()) {
            change.getJavaFunctionsUpdate().forEach(f -> javaFunctionsMerge.put(f.getName(), f));
        }

        if (change.hasCompositeFunctionsUpdate()) {
            change.getCompositeFunctionsUpdate().forEach(f -> compositeFunctionsMerge.put(f.getName(), f));
        }

        if (change.hasGroovyFunctionsUpdate()) {
            change.getGroovyFunctionsUpdate().forEach(f -> groovyFunctionsMerge.put(f.getName(), f));
        }

        if (change.hasPythonFunctionsUpdate()) {
            change.getPythonFunctionsUpdate().forEach(f -> pythonFunctionsMerge.put(f.getName(), f));
        }

        if (change.hasRulesUpdate()) {
            change.getRulesUpdate().forEach(r -> rulesMerge.put(r.getName(), r));
        }

        if (change.hasSetsUpdate()) {
            change.getSetsUpdate().forEach(s -> setsMerge.put(s.getName(), s));
        }

        if (change.hasAssignmentsUpdate()) {
            change.getAssignmentsUpdate().forEach(a -> assignmentsMerge.put(a.getNameSpace(), a));
        }

        // 3. Deletes
        if (change.hasFunctionsDelete()) {
            change.getFunctionsDelete().forEach(fn -> {
                javaFunctionsMerge.remove(fn);
                compositeFunctionsMerge.remove(fn);
                groovyFunctionsMerge.remove(fn);
                pythonFunctionsMerge.remove(fn);
            });
        }

        if (change.hasRulesDelete()) {
            change.getRulesDelete().forEach(rulesMerge::remove);
        }

        if (change.hasSetsDelete()) {
            change.getSetsDelete().forEach(setsMerge::remove);
        }

        if (change.hasAssignmentsDelete()) {
            change.getAssignmentsDelete().forEach(assignmentsMerge::remove);
        }

        // 4. Groups
        // Now change groups without pre-validation by just overwriting
        // existing groups. This has the effect, that validations will fail on publication.
        CleanseFunctionGroup root = source.getFunctionGroup();
        if (change.hasFunctionGroupUpdate()) {

            Map<String, AbstractCleanseFunctionSource<?>> allFunctions = new HashMap<>();
            allFunctions.putAll(javaFunctionsMerge);
            allFunctions.putAll(groovyFunctionsMerge);
            allFunctions.putAll(pythonFunctionsMerge);
            allFunctions.putAll(compositeFunctionsMerge);

            root = merge(change.getGroupsUpdate(), allFunctions);
        }

        // 4. Set
        target
            .withJavaFunctions(javaFunctionsMerge.values())
            .withCompositeFunctions(compositeFunctionsMerge.values())
            .withPythonFunctions(pythonFunctionsMerge.values())
            .withGroovyFunctions(groovyFunctionsMerge.values())
            .withRules(rulesMerge.values())
            .withSets(setsMerge.values())
            .withAssignments(assignmentsMerge.values())
            .withFunctionGroup(root);
    }

    private CleanseFunctionGroup merge(CleanseFunctionGroup update, Map<String, AbstractCleanseFunctionSource<?>> allFunctionsMerge) {

        Map<String, CleanseFunctionGroup> groups = updateGroupNameSet(update, StringUtils.EMPTY, new HashMap<>());
        for (Entry<String, CleanseFunctionGroup> entry : groups.entrySet()) {

            CleanseFunctionGroup each = entry.getValue();

            // Possibly update the group name in function objects.
            each.getMappedFunctions().stream()
                .map(allFunctionsMerge::get)
                .filter(Objects::nonNull)
                .forEach(fn -> fn.setGroupName(entry.getKey()));
        }

        allFunctionsMerge.values().stream()
            .filter(x -> !groups.containsKey(x.getGroupName()))
            .forEach(x -> x.setGroupName(DQUtils.ROOT_FUNCTION_GROUP_NAME));

        return update;
    }

    private Map<String, CleanseFunctionGroup> updateGroupNameSet(CleanseFunctionGroup update, String prefix, Map<String, CleanseFunctionGroup> collector) {

        String key = StringUtils.isBlank(prefix) ? update.getName() : prefix + '.' + update.getName();
        collector.put(key, update);

        for (CleanseFunctionGroup each : update.getGroups()) {
            updateGroupNameSet(each, key, collector);
        }

        return collector;
    }

    private DataQualityInstance load(String storageId) {

        DataQualityPO po = dataQualityDAO.current(storageId);
        if (Objects.isNull(po) || ArrayUtils.isEmpty(po.getContent())) {
            // Default instance
            return new DataQualityInstanceImpl(
                    DQUtils.defaultModel(storageId),
                    cleanseFunctionCacheComponent,
                    upathService);
        }

        return new DataQualityInstanceImpl(
                DataQualitySerializer.modelFromCompressedXml(po.getContent()),
                cleanseFunctionCacheComponent,
                upathService);
    }

    private void processFunctions(DataQualityInstance i, GetQualityModelContext ctx, GetQualityModelResult dto) {

        if (!ctx.isAllCleansFunctions() && CollectionUtils.isEmpty(ctx.getCleanseFunctionIds())) {
            return;
        }

        List<AbstractCleanseFunctionSource<?>> functions;
        if (ctx.isAllCleansFunctions()) {
            functions = i.getFunctions().stream()
                    .map(el -> ((AbstractCleanseFunctionImpl<?>) el).getSource())
                    .collect(Collectors.toList());
        } else {
            functions = ctx.getCleanseFunctionIds().stream()
                    .map(i::getFunction)
                    .filter(Objects::nonNull)
                    .map(el -> ((AbstractCleanseFunctionImpl<?>) el).getSource())
                    .collect(Collectors.toList());
        }

        dto.setFunctions(new GetFunctionsModelResult(functions));
    }

    private void processFunctionGroups(DataQualityInstance i, GetQualityModelContext ctx, GetQualityModelResult dto) {

        if (!ctx.isAllFunctionGroups() && CollectionUtils.isEmpty(ctx.getFunctionGroupIds())) {
            return;
        }

        Map<String, CleanseFunctionGroup> groups = new HashMap<>();
        Map<CleanseFunctionGroup, List<AbstractCleanseFunctionSource<?>>> sets = new HashMap<>();

        if (ctx.isAllFunctionGroups()) {
            i.getGroups().forEach(g -> {
                CleanseFunctionGroup cfg = ((FunctionGroupImpl) g).getSource();
                groups.put(g.getPath(), cfg);
                sets.put(cfg, g.getFunctions().stream()
                        .map(el -> ((AbstractCleanseFunctionImpl<?>) el).getSource())
                        .collect(Collectors.toList()));
            });
        } else {
            ctx.getFunctionGroupIds().stream()
                .map(i::getGroup)
                .filter(Objects::nonNull)
                .forEach(g -> {
                    CleanseFunctionGroup cfg = ((FunctionGroupImpl) g).getSource();
                    groups.put(g.getPath(), cfg);
                    sets.put(cfg, g.getFunctions().stream()
                            .map(el -> ((AbstractCleanseFunctionImpl<?>) el).getSource())
                            .collect(Collectors.toList()));
                });
        }

        dto.setFunctionGroups(new GetGroupsModelResult(groups, sets));
    }

    private void processRules(DataQualityInstance i, GetQualityModelContext ctx, GetQualityModelResult dto) {

        if (!ctx.isAllQualityRules() && CollectionUtils.isEmpty(ctx.getQualityRuleIds())) {
            return;
        }

        List<QualityRuleSource> rules;
        if (ctx.isAllQualityRules()) {
            rules = i.getRules().stream()
                    .map(el -> ((QualityRuleImpl) el).getSource())
                    .collect(Collectors.toList());
        } else {
            rules = ctx.getQualityRuleIds().stream()
                    .map(i::getRule)
                    .filter(Objects::nonNull)
                    .map(el -> ((QualityRuleImpl) el).getSource())
                    .collect(Collectors.toList());
        }

        dto.setRules(new GetRulesModelResult(rules));
    }

    private void processSets(DataQualityInstance i, GetQualityModelContext ctx, GetQualityModelResult dto) {

        if (!ctx.isAllMappingSets() && CollectionUtils.isEmpty(ctx.getMappingSetIds())) {
            return;
        }

        List<MappingSetSource> sets;
        if (ctx.isAllMappingSets()) {
            sets = i.getSets().stream()
                    .map(el -> ((MappingSetImpl) el).getSource())
                    .collect(Collectors.toList());
        } else {
            sets = ctx.getMappingSetIds().stream()
                    .map(i::getSet)
                    .filter(Objects::nonNull)
                    .map(el -> ((MappingSetImpl) el).getSource())
                    .collect(Collectors.toList());
        }

        dto.setSets(new GetSetsModelResult(sets));
    }

    private void processAssignments(DataQualityInstance i, GetQualityModelContext ctx, GetQualityModelResult dto) {

        if (!ctx.isAllAssignments() && CollectionUtils.isEmpty(ctx.getAssignmentIds())) {
            return;
        }

        List<NameSpaceAssignmentSource> assignments;
        if (ctx.isAllAssignments()) {
            assignments = i.getAssignments().stream()
                    .map(el -> ((NamespaceAssignmentImpl) el).getSource())
                    .collect(Collectors.toList());
        } else {
            assignments = ctx.getAssignmentIds().stream()
                    .map(i::getAssignment)
                    .filter(Objects::nonNull)
                    .map(el -> ((NamespaceAssignmentImpl) el).getSource())
                    .collect(Collectors.toList());
        }

        dto.setAssignments(new GetAssignmentsModelResult(assignments));
    }
}
